import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from genetic_algorithm import *
from data_loader import load_data


def visualize_solution(solution, instance):
    plt.figure(figsize=(14, 10))

    # 绘制客户点
    for c in instance.customers[1:]:
        plt.plot(c.x, c.y, 'o', markersize=8)
        plt.text(c.x, c.y, f'{c.id}', fontsize=9, ha='center', va='bottom')

    # 绘制仓库
    depot = instance.depot
    plt.plot(depot.x, depot.y, 's', markersize=12, color='r')

    # 绘制路径
    colors = plt.cm.tab10.colors
    route_details = []
    for i, route in enumerate(solution.routes):
        color = colors[i % len(colors)]
        route_x = [instance.customers[n].x for n in route]
        route_y = [instance.customers[n].y for n in route]

        plt.plot(route_x, route_y, '--', color=color, linewidth=1.5, alpha=0.8)
        for j in range(len(route) - 1):
            start = (route_x[j], route_y[j])
            end = (route_x[j + 1], route_y[j + 1])
            plt.annotate('', xy=end, xytext=start,
                         arrowprops=dict(arrowstyle="->", color=color, lw=1, alpha=0.8))

        route_details.append({
            "vehicle": i + 1,
            "customers": route[1:-1],
            "color": color
        })

    # 创建图例
    legend_elements = []
    for detail in route_details:
        legend_elements.append(Line2D([0], [0],
                                      color=detail['color'],
                                      lw=2,
                                      label=f'Vehicle {detail["vehicle"]}: {detail["customers"]}'))

    plt.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(1, 1))
    plt.title(f"VRPTW Solution (Total Cost: {solution.cost:.2f}¥)")
    plt.xlabel("X Coordinate")
    plt.ylabel("Y Coordinate")
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()


def main():
    instance = load_data()

    ga = EnhancedGeneticAlgorithm(
        instance,
        pop_size=200,
        min_pop_size=30,
        max_generations=100,
        crossover_rate=0.8,
        mutation_rate=0.1,
        T=10
    )

    best_solution, history = ga.solve()

    print("\nOptimal Solution Details:")
    for i, route in enumerate(best_solution.routes):
        print(f"Vehicle {i + 1} Route: {route}")

    visualize_solution(best_solution, instance)


if __name__ == "__main__":
    main()
